import cv2
import numpy as np

# ======================
# Step 1: 获取用户手动点击的四个点
# ======================
def get_four_points(image):
    points = []

    def mouse_handler(event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN and len(points) < 4:
            points.append((x, y))
            print(f"Point {len(points)}: ({x}, {y})")
            cv2.circle(image, (x, y), 5, (0, 255, 0), -1)
            cv2.imshow("Select 4 Points", image)

    cv2.imshow("Select 4 Points", image)
    cv2.setMouseCallback("Select 4 Points", mouse_handler)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    return points

# ======================
# Step 2: 视频俯视变换函数
# ======================
def convert_to_topview(input_video_path, output_video_path, output_size=(640, 640)):
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        print("无法打开视频文件")
        return

    # 读取第一帧选择转换区域
    ret, frame = cap.read()
    if not ret:
        print("无法读取视频第一帧")
        return

    frame_copy = frame.copy()
    print("请点击图像中的四个角点，按顺时针或逆时针顺序选择目标平面区域")
    src_pts = get_four_points(frame_copy)

    if len(src_pts) != 4:
        print("必须选择四个点！")
        return

    # 定义目标点（俯视图矩形）
    dst_pts = np.array([
        [0, 0],
        [output_size[0]-1, 0],
        [output_size[0]-1, output_size[1]-1],
        [0, output_size[1]-1]
    ], dtype="float32")

    # 计算 Homography 矩阵
    src_pts_np = np.array(src_pts, dtype="float32")
    H = cv2.getPerspectiveTransform(src_pts_np, dst_pts)

    # 初始化视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    fps = cap.get(cv2.CAP_PROP_FPS)
    out = cv2.VideoWriter(output_video_path, fourcc, fps, output_size)

    print("开始处理每一帧...")
    frame_count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        warped = cv2.warpPerspective(frame, H, output_size)
        out.write(warped)
        frame_count += 1

        if frame_count % 30 == 0:
            print(f"处理帧数: {frame_count}")

    cap.release()
    out.release()
    print("处理完成，结果保存为：", output_video_path)

# ======================
# Step 3: 主程序入口
# ======================
if __name__ == "__main__":
    input_video = "video4.mp4"
    output_video = "topview4.mp4"
    convert_to_topview(input_video, output_video)
